#!/usr/bin/env python3
"""
实时数据获取模块
支持多个数据源的实时行情数据获取
"""

import time
import json
import requests
import pandas as pd
from datetime import datetime, timedelta
from typing import Dict, List, Tuple, Optional, Any
from concurrent.futures import ThreadPoolExecutor, as_completed
import logging

from src.utils.logger import get_logger

logger = get_logger("RealTimeData")


class RealTimeDataFetcher:
    """实时数据获取器"""
    
    def __init__(self):
        self.logger = logger
        self.session = requests.Session()
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
        })
        
        # 数据源配置
        self.sources = {
            'sina': self._fetch_sina_realtime,
            'tencent': self._fetch_tencent_realtime,
            'eastmoney': self._fetch_eastmoney_realtime,
        }
        
    def get_realtime_quote(self, symbol: str, source: str = 'auto') -> Dict[str, Any]:
        """
        获取单个股票的实时行情
        
        Args:
            symbol: 股票代码
            source: 数据源，默认auto自动选择
            
        Returns:
            实时行情数据字典
        """
        try:
            if source == 'auto':
                # 自动选择可用的数据源
                for source_name, fetch_func in self.sources.items():
                    try:
                        data = fetch_func(symbol)
                        if data:
                            data['source'] = source_name
                            return data
                    except Exception as e:
                        self.logger.debug(f"{source_name} failed for {symbol}: {e}")
                        continue
            else:
                # 使用指定的数据源
                if source in self.sources:
                    data = self.sources[source](symbol)
                    if data:
                        data['source'] = source
                    return data
                    
            return self._empty_quote(symbol)
            
        except Exception as e:
            self.logger.error(f"Failed to get realtime quote for {symbol}: {e}")
            return self._empty_quote(symbol)
            
    def get_batch_quotes(self, symbols: List[str], source: str = 'auto', 
                        max_workers: int = 5) -> Dict[str, Dict]:
        """
        批量获取实时行情
        
        Args:
            symbols: 股票代码列表
            source: 数据源
            max_workers: 并发线程数
            
        Returns:
            {symbol: quote_data} 字典
        """
        results = {}
        
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_symbol = {
                executor.submit(self.get_realtime_quote, symbol, source): symbol
                for symbol in symbols
            }
            
            for future in as_completed(future_to_symbol):
                symbol = future_to_symbol[future]
                try:
                    data = future.result()
                    results[symbol] = data
                except Exception as e:
                    self.logger.error(f"Failed to get quote for {symbol}: {e}")
                    results[symbol] = self._empty_quote(symbol)
                    
        return results
        
    def _fetch_sina_realtime(self, symbol: str) -> Optional[Dict]:
        """从新浪获取实时数据"""
        try:
            # 构建新浪接口URL
            if symbol.startswith('6'):
                sina_code = f'sh{symbol}'
            elif symbol.startswith('0') or symbol.startswith('3'):
                sina_code = f'sz{symbol}'
            else:
                return None
                
            url = f'http://hq.sinajs.cn/list={sina_code}'
            
            response = self.session.get(url, timeout=5)
            response.encoding = 'gbk'
            
            if response.status_code == 200:
                content = response.text
                if 'var hq_str_' in content:
                    # 解析数据
                    data_str = content.split('="')[1].split('";')[0]
                    if not data_str:
                        return None
                        
                    fields = data_str.split(',')
                    if len(fields) < 32:
                        return None
                        
                    return {
                        'symbol': symbol,
                        'name': fields[0],
                        'open': float(fields[1]) if fields[1] else 0,
                        'prev_close': float(fields[2]) if fields[2] else 0,
                        'price': float(fields[3]) if fields[3] else 0,
                        'high': float(fields[4]) if fields[4] else 0,
                        'low': float(fields[5]) if fields[5] else 0,
                        'bid': float(fields[6]) if fields[6] else 0,
                        'ask': float(fields[7]) if fields[7] else 0,
                        'volume': int(fields[8]) if fields[8] else 0,
                        'amount': float(fields[9]) if fields[9] else 0,
                        'bid_volume': {
                            '1': int(fields[10]) if fields[10] else 0,
                            '2': int(fields[12]) if fields[12] else 0,
                            '3': int(fields[14]) if fields[14] else 0,
                            '4': int(fields[16]) if fields[16] else 0,
                            '5': int(fields[18]) if fields[18] else 0,
                        },
                        'bid_price': {
                            '1': float(fields[11]) if fields[11] else 0,
                            '2': float(fields[13]) if fields[13] else 0,
                            '3': float(fields[15]) if fields[15] else 0,
                            '4': float(fields[17]) if fields[17] else 0,
                            '5': float(fields[19]) if fields[19] else 0,
                        },
                        'ask_volume': {
                            '1': int(fields[20]) if fields[20] else 0,
                            '2': int(fields[22]) if fields[22] else 0,
                            '3': int(fields[24]) if fields[24] else 0,
                            '4': int(fields[26]) if fields[26] else 0,
                            '5': int(fields[28]) if fields[28] else 0,
                        },
                        'ask_price': {
                            '1': float(fields[21]) if fields[21] else 0,
                            '2': float(fields[23]) if fields[23] else 0,
                            '3': float(fields[25]) if fields[25] else 0,
                            '4': float(fields[27]) if fields[27] else 0,
                            '5': float(fields[29]) if fields[29] else 0,
                        },
                        'date': fields[30] if len(fields) > 30 else '',
                        'time': fields[31] if len(fields) > 31 else '',
                        'status': 'open' if fields[3] != '0' else 'closed',
                        'timestamp': datetime.now().isoformat(),
                    }
                    
        except Exception as e:
            self.logger.debug(f"Sina fetch error for {symbol}: {e}")
            
        return None
        
    def _fetch_tencent_realtime(self, symbol: str) -> Optional[Dict]:
        """从腾讯获取实时数据"""
        try:
            # 构建腾讯接口URL
            if symbol.startswith('6'):
                qq_code = f'sh{symbol}'
            elif symbol.startswith('0') or symbol.startswith('3'):
                qq_code = f'sz{symbol}'
            else:
                return None
                
            url = f'http://qt.gtimg.cn/q={qq_code}'
            
            response = self.session.get(url, timeout=5)
            response.encoding = 'gbk'
            
            if response.status_code == 200 and 'v_' in response.text:
                # 解析数据
                data_str = response.text.split('~')
                if len(data_str) > 45:
                    return {
                        'symbol': symbol,
                        'name': data_str[1],
                        'price': float(data_str[3]) if data_str[3] else 0,
                        'prev_close': float(data_str[4]) if data_str[4] else 0,
                        'open': float(data_str[5]) if data_str[5] else 0,
                        'volume': int(data_str[6]) * 100 if data_str[6] else 0,  # 手数转股数
                        'bid': float(data_str[9]) if data_str[9] else 0,
                        'ask': float(data_str[19]) if data_str[19] else 0,
                        'high': float(data_str[33]) if data_str[33] else 0,
                        'low': float(data_str[34]) if data_str[34] else 0,
                        'amount': float(data_str[37]) * 10000 if data_str[37] else 0,  # 万元转元
                        'change': float(data_str[31]) if data_str[31] else 0,
                        'change_percent': float(data_str[32]) if data_str[32] else 0,
                        'pe_ratio': float(data_str[39]) if data_str[39] and data_str[39] != '' else 0,
                        'pb_ratio': float(data_str[46]) if len(data_str) > 46 and data_str[46] else 0,
                        'market_cap': float(data_str[45]) * 100000000 if len(data_str) > 45 and data_str[45] else 0,  # 亿转元
                        'turnover_rate': float(data_str[38]) if data_str[38] else 0,
                        'timestamp': datetime.now().isoformat(),
                    }
                    
        except Exception as e:
            self.logger.debug(f"Tencent fetch error for {symbol}: {e}")
            
        return None
        
    def _fetch_eastmoney_realtime(self, symbol: str) -> Optional[Dict]:
        """从东方财富获取实时数据"""
        try:
            # 构建东方财富接口URL
            if symbol.startswith('6'):
                em_code = f'1.{symbol}'
            elif symbol.startswith('0') or symbol.startswith('3'):
                em_code = f'0.{symbol}'
            else:
                return None
                
            url = f'http://push2.eastmoney.com/api/qt/stock/get'
            params = {
                'secid': em_code,
                'fields': 'f43,f44,f45,f46,f47,f48,f49,f50,f51,f52,f57,f58,f59,f60,f61,f62,f63,f64,f65,f66,f67,f68,f69,f70,f71,f116,f117,f162,f163,f164,f165,f166,f167,f168,f169,f170,f171',
                'invt': '2',
                'fltt': '2',
                '_': str(int(time.time() * 1000))
            }
            
            response = self.session.get(url, params=params, timeout=5)
            
            if response.status_code == 200:
                data = response.json()
                if data and 'data' in data and data['data']:
                    stock = data['data']
                    return {
                        'symbol': symbol,
                        'name': stock.get('f58', ''),
                        'price': stock.get('f43', 0) / 100 if stock.get('f43') else 0,
                        'prev_close': stock.get('f60', 0) / 100 if stock.get('f60') else 0,
                        'open': stock.get('f46', 0) / 100 if stock.get('f46') else 0,
                        'high': stock.get('f44', 0) / 100 if stock.get('f44') else 0,
                        'low': stock.get('f45', 0) / 100 if stock.get('f45') else 0,
                        'volume': stock.get('f47', 0),
                        'amount': stock.get('f48', 0),
                        'change': stock.get('f169', 0) / 100 if stock.get('f169') else 0,
                        'change_percent': stock.get('f170', 0) / 100 if stock.get('f170') else 0,
                        'amplitude': stock.get('f171', 0) / 100 if stock.get('f171') else 0,
                        'turnover_rate': stock.get('f168', 0) / 100 if stock.get('f168') else 0,
                        'pe_ratio': stock.get('f162', 0) / 100 if stock.get('f162') else 0,
                        'pb_ratio': stock.get('f167', 0) / 100 if stock.get('f167') else 0,
                        'market_cap': stock.get('f116', 0),
                        'circulating_cap': stock.get('f117', 0),
                        'timestamp': datetime.now().isoformat(),
                    }
                    
        except Exception as e:
            self.logger.debug(f"EastMoney fetch error for {symbol}: {e}")
            
        return None
        
    def _empty_quote(self, symbol: str) -> Dict:
        """返回空的行情数据"""
        return {
            'symbol': symbol,
            'name': '',
            'price': 0,
            'prev_close': 0,
            'open': 0,
            'high': 0,
            'low': 0,
            'volume': 0,
            'amount': 0,
            'change': 0,
            'change_percent': 0,
            'status': 'error',
            'timestamp': datetime.now().isoformat(),
        }
        
    def get_market_overview(self) -> Dict[str, Any]:
        """获取市场概览数据"""
        try:
            # 获取主要指数
            indices = {
                '000001': '上证指数',
                '399001': '深证成指',
                '399006': '创业板指',
                '000300': '沪深300',
                '000016': '上证50',
                '000905': '中证500',
            }
            
            index_data = self.get_batch_quotes(list(indices.keys()))
            
            # 获取涨跌统计
            stats = self._get_market_stats()
            
            return {
                'indices': index_data,
                'stats': stats,
                'timestamp': datetime.now().isoformat(),
            }
            
        except Exception as e:
            self.logger.error(f"Failed to get market overview: {e}")
            return {}
            
    def _get_market_stats(self) -> Dict[str, int]:
        """获取市场涨跌统计"""
        try:
            # 这里可以调用东方财富的统计接口
            url = 'http://push2.eastmoney.com/api/qt/ulist.np/get'
            params = {
                'fltt': '2',
                'fields': 'f3',
                'secids': '0.000001,1.000001',
                '_': str(int(time.time() * 1000))
            }
            
            response = self.session.get(url, params=params, timeout=5)
            
            if response.status_code == 200:
                data = response.json()
                # 解析统计数据
                # 这里简化处理，实际需要更复杂的逻辑
                return {
                    'up': 0,
                    'down': 0,
                    'flat': 0,
                    'halt': 0,
                }
                
        except Exception as e:
            self.logger.debug(f"Failed to get market stats: {e}")
            
        return {
            'up': 0,
            'down': 0,
            'flat': 0,
            'halt': 0,
        }
        

class WebSocketDataStream:
    """WebSocket实时数据流（用于推送实时数据）"""
    
    def __init__(self, socketio):
        self.socketio = socketio
        self.fetcher = RealTimeDataFetcher()
        self.subscriptions = {}  # {client_id: [symbols]}
        self.running = False
        
    def subscribe(self, client_id: str, symbols: List[str]):
        """订阅股票实时数据"""
        self.subscriptions[client_id] = symbols
        self.logger.info(f"Client {client_id} subscribed to {symbols}")
        
    def unsubscribe(self, client_id: str):
        """取消订阅"""
        if client_id in self.subscriptions:
            del self.subscriptions[client_id]
            self.logger.info(f"Client {client_id} unsubscribed")
            
    def start_streaming(self):
        """开始推送实时数据"""
        self.running = True
        
        while self.running:
            try:
                # 获取所有订阅的股票
                all_symbols = set()
                for symbols in self.subscriptions.values():
                    all_symbols.update(symbols)
                    
                if all_symbols:
                    # 批量获取实时数据
                    quotes = self.fetcher.get_batch_quotes(list(all_symbols))
                    
                    # 推送给订阅的客户端
                    for client_id, symbols in self.subscriptions.items():
                        client_quotes = {s: quotes[s] for s in symbols if s in quotes}
                        self.socketio.emit('realtime_data', client_quotes, room=client_id)
                        
                # 每3秒更新一次
                time.sleep(3)
                
            except Exception as e:
                self.logger.error(f"Streaming error: {e}")
                time.sleep(5)
                
    def stop_streaming(self):
        """停止推送"""
        self.running = False